import numpy as np
from scipy import special
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread


class Form(HasTraits):
    frequencies = List(Float, [3, 10])
    elevations = List(Float, [0, 10])
    azimuthStart = Float(0)
    azimuthStop = Float(45)
    azimuthNumber = Int(91)
    plot = Instance(Plot)
    progress = Str('0%')
    thread = QThread()

    def calculate(self):
        np.seterr(divide='ignore')
        phi = np.arange(0, 361)
        s = np.column_stack((np.cos(np.deg2rad(phi)), np.sin(np.deg2rad(phi))))
        k = 2 * np.pi / .3 * .6
        v = np.arange(0, 360, 8)
        v = np.column_stack((np.cos(np.deg2rad(v)), np.sin(np.deg2rad(v))))
        n = np.cross([0, 0, 1], np.roll(v, -1, 0) - v)[:, 0:-1]
        l = np.linalg.norm(n, axis=1)
        c = (np.roll(v, -1, 0) + v) / 2
        V = -np.exp(1j * k * c[:, 0])
        r = c[np.newaxis, :, :] - c[:, np.newaxis, :]
        R = np.linalg.norm(r, axis=2)
        Z = special.hankel2(0, k * R) * l / 4j
        Z[range(l.size), range(l.size)] = (1 - 2j / np.pi * (np.log(k * l / 4) - 1 + np.euler_gamma)) * l / 4j
        atm = np.linalg.solve(Z, V)
        rcstm = 2 * np.exp(1j * np.pi / 4) / np.sqrt(k) * np.exp(1j * k * s @ c.T) * l / 4j @ atm / np.pi ** .5
        Z = k * (n[np.newaxis, :, :] * r).sum(2) / R * special.hankel2(1, k * R) / 4j
        Z[range(l.size), range(l.size)] = -.5
        ate = np.linalg.solve(Z, V)
        rcste = -2j * np.exp(1j * np.pi / 4) * np.sqrt(k) * s @ n.T * np.exp(1j * k * s @ c.T) / 4j @ ate / np.sqrt(np.pi)
        n = np.arange(-80, 81)
        ka = k
        jn = special.jv(n, ka)
        hn = special.hankel2(n, ka)
        jp = special.jvp(n, ka)
        hp = special.h2vp(n, ka)
        rcstm0 = -2 * np.exp(1j * np.pi / 4) / np.sqrt(np.pi * ka) * ((-1.) ** n) * jn / hn @ np.exp(1j * n[:, np.newaxis] * np.deg2rad(phi))
        rcste0 = -2 * np.exp(1j * np.pi / 4) / np.sqrt(np.pi * ka) * ((-1.) ** n) * jp / hp @ np.exp(1j * n[:, np.newaxis] * np.deg2rad(phi))

        self.plot.data.update_data(
            x=phi,
            tm=20*np.log10(np.abs(rcstm0)),
            te=20*np.log10(np.abs(rcste0)),
        )

    def traits_init(self):
        self.thread.run = self.calculate
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            x=[],
            tm=[],
            te=[],
        )
        self.plot.x_axis.title = "Azimuth/deg"
        self.plot.y_axis.title = "RCS/dBsm"
        self.plot.plot(("x", "tm"), name='TM', color='auto')
        self.plot.plot(("x", "te"), name='TE', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
    button=Button('计算')
    def _button_fired(self):
        self.calculate()
    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            UItem('button'),
            UItem('progress'),
            spring
        ),
        width=1200,
        height=900,
        resizable=True,
    )
    parameter_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        Group(
            UItem('elevations'),
            label='高低角/°',
            show_border=True,
        ),
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
        title='MoM 2D',
    )


if __name__ == '__main__':
    form = Form()
    form.configure_traits()

